package main

import (
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/go-playground/validator/v10"
	"github.com/sirupsen/logrus"
	"gopkg.in/yaml.v3"
)

// Config 结构体用于存储配置信息
type Config struct {
	API struct {
		SecretKey string `mapstructure:"secret_key"`
	} `mapstructure:"api"`
	Server struct {
		Host string `mapstructure:"host"`
		Port int    `mapstructure:"port"`
	} `mapstructure:"server"`
	Model struct {
		ModelPath  string `mapstructure:"model_path"`
		MaxLength  int    `mapstructure:"max_length"`
		BatchSize  int    `mapstructure:"batch_size"`
		Device     string `mapstructure:"device"`
	} `mapstructure:"model"`
	Logging struct {
		Level    string `mapstructure:"level"`
		FilePath string `mapstructure:"file_path"`
	} `mapstructure:"logging"`
}

// EmbeddingRequest 请求结构体
type EmbeddingRequest struct {
	Texts []string `json:"texts" binding:"required,min=1,max=100"`
}

// EmbeddingResponse 响应结构体
type EmbeddingResponse struct {
	Embeddings [][]float32 `json:"embeddings"`
}

// HealthResponse 健康检查响应结构体
type HealthResponse struct {
	Status     string `json:"status"`
	ModelLoaded bool  `json:"model_loaded"`
}

// RootResponse 根路径响应结构体
type RootResponse struct {
	Service    string                 `json:"service"`
	Version    string                 `json:"version"`
	Descrition string                 `json:"description"`
	Endpoints  map[string]string      `json:"endpoints"`
}

// BGEModel 定义嵌入向量模型接口
type BGEModel interface {
	Encode(texts []string) ([][]float32, error)
	Load(modelPath string) error
	IsLoaded() bool
}

// BGEModelImpl BGE模型实现
type BGEModelImpl struct {
	modelPath string
	isLoaded  bool
	// 这里可以添加实际的模型相关字段
}

// NewBGEModel 创建新的BGE模型实例
func NewBGEModel() BGEModel {
	return &BGEModelImpl{}
}

// Load 加载模型
func (m *BGEModelImpl) Load(modelPath string) error {
	// 这里应该是实际的模型加载逻辑
	// 由于Go中没有直接对应的transformers库，这里需要通过外部调用或其他方式实现
	// 可以考虑使用ONNX Runtime、TorchScript或调用Python子进程等方式
	
	// 检查模型路径是否存在，如果不存在则创建一个
	if _, err := os.Stat(modelPath); os.IsNotExist(err) {
		logger.Infof("Model path does not exist, creating directory: %s", modelPath)
		err := os.MkdirAll(modelPath, 0755)
		if err != nil {
			logger.Warnf("Failed to create model directory: %v", err)
		}
	}
	
	m.modelPath = modelPath
	m.isLoaded = true
	
	return nil
}

// IsLoaded 检查模型是否已加载
func (m *BGEModelImpl) IsLoaded() bool {
	return m.isLoaded
}

// Encode 编码文本为嵌入向量
func (m *BGEModelImpl) Encode(texts []string) ([][]float32, error) {
	if !m.isLoaded {
		return nil, fmt.Errorf("model is not loaded")
	}
	
	// 这里应该是实际的文本编码逻辑
	// 由于Go中没有直接对应的transformers库，这里需要通过外部调用或其他方式实现
	// 以下是模拟实现，返回随机向量
	
	// 创建一个与输入文本数量相同的二维切片
	embeddings := make([][]float32, len(texts))
	for i := range embeddings {
		// 每个文本生成一个768维的向量
		embedding := make([]float32, 768)
		// 在实际应用中，这里应该是真实的嵌入向量
		// 这里只是示例，填充随机值
		for j := range embedding {
			embedding[j] = 0.1 // 示例值
		}
		embeddings[i] = embedding
	}
	
	return embeddings, nil
}

// 全局变量
var (
	config Config
	model  BGEModel
	logger *logrus.Logger
	validate = validator.New()
)

// initConfig 初始化配置
func initConfig() {
	// 设置配置文件路径
	configPath := filepath.Join(".", "config.yaml")
	
	// 读取配置文件
	data, err := os.ReadFile(configPath)
	if err != nil {
		log.Fatalf("Failed to read config file: %v", err)
	}
	
	// 解析YAML配置
	err = yaml.Unmarshal(data, &config)
	if err != nil {
		log.Fatalf("Failed to parse config file: %v", err)
	}
	
	// 设置默认值
	if config.API.SecretKey == "" {
		config.API.SecretKey = "bge-service-secret-key-2023"
	}
	if config.Server.Host == "" {
		config.Server.Host = "0.0.0.0"
	}
	if config.Server.Port == 0 {
		config.Server.Port = 8000
	}
	if config.Model.ModelPath == "" {
		config.Model.ModelPath = "./models/bge-base-zh"
	}
	if config.Model.MaxLength == 0 {
		config.Model.MaxLength = 512
	}
	if config.Model.BatchSize == 0 {
		config.Model.BatchSize = 32
	}
	if config.Model.Device == "" {
		config.Model.Device = "cpu"
	}
	if config.Logging.Level == "" {
		config.Logging.Level = "INFO"
	}
}

// initLogger 初始化日志
func initLogger() {
	logger = logrus.New()
	
	// 设置日志级别
	level, err := logrus.ParseLevel(config.Logging.Level)
	if err != nil {
		level = logrus.InfoLevel
	}
	logger.SetLevel(level)
	
	// 设置日志格式
	logger.SetFormatter(&logrus.JSONFormatter{
		TimestampFormat: "2006-01-02 15:04:05",
	})
	
	// 设置日志输出
	if config.Logging.FilePath != "" {
		file, err := os.OpenFile(config.Logging.FilePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
		if err == nil {
			logger.SetOutput(file)
		} else {
			logger.SetOutput(os.Stdout)
			logger.Warnf("Failed to open log file: %v, using stdout instead", err)
		}
	} else {
		logger.SetOutput(os.Stdout)
	}
}

// initModel 初始化模型
func initModel() {
	// 使用Python后端实现，调用真实的BGE模型
	factory := ModelFactory{}
	model = factory.NewModel("python")
	
	logger.Infof("Loading model from: %s", config.Model.ModelPath)
	
	err := model.Load(config.Model.ModelPath)
	if err != nil {
		logger.Fatalf("Failed to load model: %v", err)
	}
	
	logger.Info("Model loaded successfully")
}

// verifyAPIKey 验证API密钥
func verifyAPIKey(c *gin.Context) {
	authHeader := c.GetHeader("Authorization")
	if authHeader == "" {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Authorization header is missing"})
		return
	}
	
	// 检查Authorization头格式
	if !strings.HasPrefix(authHeader, "Bearer ") {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid authorization format. Expected 'Bearer <token>'"})
		return
	}
	
	// 提取token
	tokenString := strings.TrimPrefix(authHeader, "Bearer ")
	
	// 直接比较token与配置中的密钥
	if tokenString != config.API.SecretKey {
		c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid API key"})
		return
	}
	
	c.Next()
}

// getEmbeddingsHandler 处理嵌入向量请求
func getEmbeddingsHandler(c *gin.Context) {
	var request EmbeddingRequest
	
	// 绑定JSON请求体
	if err := c.ShouldBindJSON(&request); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	
	// 验证输入
	if len(request.Texts) == 0 {
		c.JSON(http.StatusBadRequest, gin.H{"error": "No text provided for embedding"})
		return
	}
	
	logger.Infof("Processing %d texts for embedding", len(request.Texts))
	
	// 获取嵌入向量
	embeddings, err := model.Encode(request.Texts)
	if err != nil {
		logger.Errorf("Error generating embeddings: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("Internal server error: %s", err.Error())})
		return
	}
	
	logger.Infof("Successfully generated %d embeddings with dimension: %d", len(embeddings), len(embeddings[0]))
	
	c.JSON(http.StatusOK, EmbeddingResponse{Embeddings: embeddings})
}

// rootHandler 根路径处理器
func rootHandler(c *gin.Context) {
	response := RootResponse{
		Service:    "BGE Embedding Service",
		Version:    "1.0.0",
		Descrition: "A service for generating text embeddings using BGE models",
		Endpoints: map[string]string{
			"embeddings": "/v1/embeddings",
			"docs":       "/docs",
		},
	}
	c.JSON(http.StatusOK, response)
}

// healthHandler 健康检查处理器
func healthHandler(c *gin.Context) {
	response := HealthResponse{
		Status:     "healthy",
		ModelLoaded: model.IsLoaded(),
	}
	c.JSON(http.StatusOK, response)
}

// setupRouter 设置路由
func setupRouter() *gin.Engine {
	router := gin.Default()
	
	// 根路径
	router.GET("/", rootHandler)
	
	// 健康检查
	router.GET("/health", healthHandler)
	
	// API路由组
	api := router.Group("/v1")
	api.Use(verifyAPIKey)
	{
		api.POST("/embeddings", getEmbeddingsHandler)
	}
	
	return router
}

func main() {
	// 初始化配置
	initConfig()
	
	// 初始化日志
	initLogger()
	
	// 初始化模型
	initModel()
	
	// 设置路由
	router := setupRouter()
	
	// 启动服务器
	addr := fmt.Sprintf("%s:%d", config.Server.Host, config.Server.Port)
	logger.Infof("Starting BGE Embedding Service on %s", addr)
	
	server := &http.Server{
		Addr:         addr,
		Handler:      router,
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 30 * time.Second,
		IdleTimeout:  60 * time.Second,
	}
	
	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logger.Fatalf("Failed to start server: %v", err)
	}
}